#!/usr/bin/env python3
"""
Download script for CogVideo models.
"""

import os
import logging
from huggingface_hub import hf_hub_download

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# --- Model Configuration ---
# The model names and files to download.
MODELS = {
    "cogvideo": {
        "repo_id": "THUDM/CogVideo",
        "files": [
            # Add the specific model files to download here
            # e.g., "pytorch_model.bin", "config.json"
        ]
    }
}

# --- Download Function ---
def download_model(repo_id, file, cache_dir):
    """Downloads a single model file from the Hugging Face Hub."""
    logger.info(f"Downloading {file} from {repo_id}...")
    try:
        hf_hub_download(
            repo_id=repo_id,
            filename=file,
            cache_dir=cache_dir,
            force_filename=file # Ensure the filename is not changed
        )
        logger.info(f"✓ Successfully downloaded {file}.")
    except Exception as e:
        logger.error(f"✗ Failed to download {file}: {e}")

# --- Main Execution ---
if __name__ == "__main__":
    cache_dir = os.environ.get("HF_HOME", "./models")
    logger.info(f"Models will be downloaded to: {cache_dir}")

    for model_name, model_info in MODELS.items():
        logger.info(f"--- Downloading {model_name} models ---")
        for file in model_info["files"]:
            download_model(model_info["repo_id"], file, cache_dir)

    logger.info("--- All model downloads complete! ---")

